{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pytorch_lightning as pl\n",
    "from tqdm.notebook import tqdm\n",
    "from sklearn.metrics import accuracy_score\n",
    "from transformers import (\n",
    "    DataProcessor,\n",
    "    InputExample,\n",
    "    BertTokenizer,\n",
    "    BertConfig,\n",
    "    BertForSequenceClassification,\n",
    "    glue_convert_examples_to_features,\n",
    ")\n",
    "\n",
    "def set_random_seed(seed):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "set_random_seed(2020)\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "CKPT_PATH = '../../ckpt/'\n",
    "TSV_PATH = '../../data/tsv/'\n",
    "PICKLE_PATH = '../../data/pickle/'\n",
    "RESULT_PATH = '../../data/result/'\n",
    "PRETRAINED_PATH = '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityLinkingProcessor(DataProcessor):\n",
    "    \"\"\"实体链接数据处理\"\"\"\n",
    "\n",
    "    def get_train_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='train',\n",
    "        )\n",
    "\n",
    "    def get_dev_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='valid',\n",
    "        )\n",
    "\n",
    "    def get_test_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='test',\n",
    "        )\n",
    "\n",
    "    def get_labels(self):\n",
    "        return ['0', '1']\n",
    "\n",
    "    def _create_examples(self, lines, set_type):\n",
    "        examples = []\n",
    "        for i, line in enumerate(lines):\n",
    "            if i == 0:\n",
    "                continue\n",
    "            guid = f'{set_type}-{i}'\n",
    "            text_a = line[1] + ' ' + line[3]\n",
    "            text_b = line[5]\n",
    "            label = line[-1]\n",
    "            examples.append(InputExample(\n",
    "                guid=guid,\n",
    "                text_a=text_a,\n",
    "                text_b=text_b,\n",
    "                label=label,\n",
    "            ))\n",
    "        return examples\n",
    "\n",
    "    def create_dataloader(self, examples, tokenizer, max_length=384,\n",
    "                          shuffle=False, batch_size=32, use_pickle=False):\n",
    "        pickle_name = 'EL_FEATURE_' + examples[0].guid.split('-')[0].upper() + '.pkl'\n",
    "        if use_pickle:\n",
    "            features = pd.read_pickle(PICKLE_PATH + pickle_name)\n",
    "        else:\n",
    "            features = glue_convert_examples_to_features(\n",
    "                examples,\n",
    "                tokenizer,\n",
    "                label_list=self.get_labels(),\n",
    "                max_length=max_length,\n",
    "                output_mode='classification',\n",
    "                pad_on_left=False,\n",
    "                pad_token=tokenizer.pad_token_id,\n",
    "                pad_token_segment_id=0,\n",
    "            )\n",
    "            pd.to_pickle(features, PICKLE_PATH + pickle_name)\n",
    "\n",
    "        dataset = torch.utils.data.TensorDataset(\n",
    "            torch.LongTensor([f.input_ids for f in features]),\n",
    "            torch.LongTensor([f.attention_mask for f in features]),\n",
    "            torch.LongTensor([f.token_type_ids for f in features]),\n",
    "            torch.LongTensor([f.label for f in features]),\n",
    "        )\n",
    "\n",
    "        dataloader = torch.utils.data.DataLoader(\n",
    "            dataset,\n",
    "            shuffle=shuffle,\n",
    "            batch_size=batch_size,\n",
    "        )\n",
    "        return dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:transformers.tokenization_utils:Model name '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/' not found in model shortcut name list (bert-base-uncased, bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, bert-base-multilingual-cased, bert-base-chinese, bert-base-german-cased, bert-large-uncased-whole-word-masking, bert-large-cased-whole-word-masking, bert-large-uncased-whole-word-masking-finetuned-squad, bert-large-cased-whole-word-masking-finetuned-squad, bert-base-cased-finetuned-mrpc, bert-base-german-dbmdz-cased, bert-base-german-dbmdz-uncased, bert-base-finnish-cased-v1, bert-base-finnish-uncased-v1, bert-base-dutch-cased). Assuming '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/' is a path, a model identifier, or url to a directory containing tokenizer files.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/added_tokens.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/special_tokens_map.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/tokenizer_config.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:loading file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/vocab.txt\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.tokenization_utils:loading file None\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[tensor([[ 101, 1921,  678,  ...,    0,    0,    0],\n",
      "        [ 101, 3719, 1649,  ...,    0,    0,    0],\n",
      "        [ 101, 3719, 1649,  ...,    0,    0,    0],\n",
      "        ...,\n",
      "        [ 101, 1322, 2791,  ...,    0,    0,    0],\n",
      "        [ 101, 1139, 4909,  ...,    0,    0,    0],\n",
      "        [ 101, 1139, 4909,  ...,    0,    0,    0]]), tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0]]), tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0]]), tensor([1, 0, 0, 1, 1, 0, 0, 1])]\n"
     ]
    }
   ],
   "source": [
    "processor = EntityLinkingProcessor()\n",
    "tokenizer = BertTokenizer.from_pretrained(PRETRAINED_PATH)\n",
    "valid_examples = processor.get_dev_examples(TSV_PATH + 'EL_VALID.tsv')\n",
    "valid_loader = processor.create_dataloader(\n",
    "    examples=valid_examples,\n",
    "    tokenizer=tokenizer,\n",
    "    max_length=384,\n",
    "    shuffle=False,\n",
    "    batch_size=8,\n",
    "    use_pickle=True,\n",
    ")\n",
    "\n",
    "for batch in valid_loader:\n",
    "    print(batch)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityLinkingModel(pl.LightningModule):\n",
    "    \"\"\"实体链接模型\"\"\"\n",
    "\n",
    "    def __init__(self, max_length=384, batch_size=32, use_pickle=True):\n",
    "        super(EntityLinkingModel, self).__init__()\n",
    "        # 输入最大长度\n",
    "        self.max_length = max_length\n",
    "        self.batch_size = batch_size\n",
    "        self.use_pickle = use_pickle\n",
    "\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(PRETRAINED_PATH)\n",
    "\n",
    "        # 预训练模型配置信息\n",
    "        self.config = BertConfig.from_json_file(PRETRAINED_PATH + 'bert_config.json')\n",
    "        self.config.num_labels = 1\n",
    "\n",
    "        # 预训练模型\n",
    "        self.bert = BertForSequenceClassification.from_pretrained(\n",
    "            PRETRAINED_PATH + 'pytorch_model.bin',\n",
    "            config=self.config,\n",
    "        )\n",
    "\n",
    "        # 二分类损失函数\n",
    "        self.criterion = nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask, token_type_ids):\n",
    "        logits = self.bert(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids,\n",
    "        )[0]\n",
    "        return logits.squeeze()\n",
    "\n",
    "    def prepare_data(self):\n",
    "        self.processor = EntityLinkingProcessor()\n",
    "        self.train_examples = self.processor.get_train_examples(TSV_PATH + 'EL_TRAIN.tsv')\n",
    "        self.valid_examples = self.processor.get_dev_examples(TSV_PATH + 'EL_VALID.tsv')\n",
    "        self.test_examples = self.processor.get_test_examples(TSV_PATH + 'EL_TEST.tsv')\n",
    "\n",
    "        self.train_loader = self.processor.create_dataloader(\n",
    "            examples=self.train_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=True,\n",
    "            batch_size=32,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "        self.valid_loader = self.processor.create_dataloader(\n",
    "            examples=self.valid_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=32,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "        self.test_loader = self.processor.create_dataloader(\n",
    "            examples=self.test_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=32,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, labels = batch\n",
    "        logits = self(input_ids, attention_mask, token_type_ids)\n",
    "        loss = self.criterion(logits, labels.float())\n",
    "\n",
    "        preds = (logits > 0).int()\n",
    "        acc = (preds == labels).float().mean()\n",
    "\n",
    "        tensorboard_logs = {'train_loss': loss, 'train_acc': acc}\n",
    "        return {'loss': loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, labels = batch\n",
    "        logits = self(input_ids, attention_mask, token_type_ids)\n",
    "        loss = self.criterion(logits, labels.float())\n",
    "\n",
    "        preds = (logits > 0).int()\n",
    "        acc = (preds == labels).float().mean()\n",
    "\n",
    "        return {'val_loss': loss, 'val_acc': acc}\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        val_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
    "        val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()\n",
    "\n",
    "        tensorboard_logs = {'val_loss': val_loss, 'val_acc': val_acc}\n",
    "        return {'val_loss': val_loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam([p for p in self.parameters() if p.requires_grad], lr=2e-5, eps=1e-8)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return self.train_loader\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.valid_loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:transformers.tokenization_utils:Model name '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/' not found in model shortcut name list (bert-base-uncased, bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, bert-base-multilingual-cased, bert-base-chinese, bert-base-german-cased, bert-large-uncased-whole-word-masking, bert-large-cased-whole-word-masking, bert-large-uncased-whole-word-masking-finetuned-squad, bert-large-cased-whole-word-masking-finetuned-squad, bert-base-cased-finetuned-mrpc, bert-base-german-dbmdz-cased, bert-base-german-dbmdz-uncased, bert-base-finnish-cased-v1, bert-base-finnish-uncased-v1, bert-base-dutch-cased). Assuming '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/' is a path, a model identifier, or url to a directory containing tokenizer files.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/added_tokens.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/special_tokens_map.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/tokenizer_config.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:loading file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/vocab.txt\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.modeling_utils:loading weights file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/pytorch_model.bin\n",
      "INFO:transformers.modeling_utils:Weights of BertForSequenceClassification not initialized from pretrained model: ['classifier.weight', 'classifier.bias']\n",
      "INFO:transformers.modeling_utils:Weights from pretrained model not used in BertForSequenceClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4719eaa53d14b438ea1d9e93e11912c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=17868.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model = EntityLinkingModel.load_from_checkpoint(\n",
    "    checkpoint_path=CKPT_PATH+'EL_BASE_EPOCH0.ckpt',\n",
    ")\n",
    "model.to(device)\n",
    "model.eval()\n",
    "\n",
    "result = []\n",
    "for batch in tqdm(valid_loader):\n",
    "    for i in range(len(batch)):\n",
    "        batch[i] = batch[i].to(device)\n",
    "    input_ids, attention_mask, token_type_ids, labels = batch\n",
    "    logits = model(input_ids, attention_mask, token_type_ids)\n",
    "    preds = (logits > 0).int()\n",
    "    result.extend(preds.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_data = pd.read_csv(TSV_PATH+'EL_VALID.tsv', sep='\\t')\n",
    "valid_data['result'] = result\n",
    "valid_data.to_csv(RESULT_PATH+'EL_VALID_RESULT.tsv', index=False, sep='\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9549318240648108\n"
     ]
    }
   ],
   "source": [
    "print(accuracy_score(valid_data['predict'], valid_data['result']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
